from torchvision import models
import torch.nn as nn
import torch
from utils.args import get_args

a = models.Swin_B_Weights
# 网络权重初始化
def weight_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
    # 也可以判断是否为conv2d，使用相应的初始化方式
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    # 是否为批归一化层
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)


def modify_fc(args, fc):
    in_channel = fc.in_features
    if args.dataset == 'cub':
        linear = nn.Linear(in_channel, 200)
    elif args.dataset == 'dogs':
        linear = nn.Linear(in_channel, 120)
    elif args.dataset == 'aircrafts':
        linear = nn.Linear(in_channel, 100)
    else:
        raise Exception(f'no such dataset named {args.dataset}')
    # 初始化全连接层权重
    nn.init.xavier_normal_(linear.weight)
    nn.init.constant_(linear.bias, 0)
    return linear


def setup(args):
    if args.model == 'resnet50':
        model = models.resnet50(weights=None)
        # 加载预训练模型
        if args.pretrained_path:
            model.load_state_dict(torch.load(args.pretrained_path))
        else:
            model.apply(weight_init)
        # 设置全连接层输出维度
        model.fc = modify_fc(args, model.fc)

    elif args.model == 'vit-b-16':
        model = models.vit_b_16(weights=None)
        # 加载预训练模型
        if args.pretrained_path:
            model.load_state_dict(torch.load(args.pretrained_path))
        else:
            model.apply(weight_init)
        # 设置输入图片大小
        model.image_size = args.image_size
        # 设置全连接层输出维度
        model.heads.head = modify_fc(args, model.heads.head)

    elif args.model == 'vit-l-16':
        model = models.vit_l_16(weights=None)
        # 加载预训练模型
        if args.pretrained_path:
            model.load_state_dict(torch.load(args.pretrained_path))
        else:
            model.apply(weight_init)
        # 设置输入图片大小
        model.image_size = args.image_size
        # 设置全连接层输出维度
        model.heads.head = modify_fc(args, model.heads.head)

    elif args.model == 'swin-t':
        model = models.swin_t(weights=None)
        # 加载预训练模型
        if args.pretrained_path:
            model.load_state_dict(torch.load(args.pretrained_path))
        else:
            model.apply(weight_init)
        # 设置输入图片大小
        model.image_size = args.image_size
        # 设置全连接层输出维度
        model.head = modify_fc(args, model.head)

    elif args.model == 'swin-b':
        model = models.swin_b(weights=None)
        # 加载预训练模型
        if args.pretrained_path:
            model.load_state_dict(torch.load(args.pretrained_path))
        else:
            model.apply(weight_init)
        # 设置输入图片大小
        model.image_size = args.image_size
        # 设置全连接层输出维度
        model.head = modify_fc(args, model.head)

    else:
        raise Exception(f'no such model names {args.model}')

    for param in model.parameters():
        param.requires_grad = True  # make parameters in model learnable

    return model


if __name__ == '__main__':
    # 获取命令行参数
    args = get_args()
    model = setup(args)
    print(model.encoder.layers[-2])
    # data = torch.randn(size=(1, 3, 448, 448))
    # pred = model(data)
